% rf_mcmc
%
% This file runs the estimation
%
%%

clear all
clc

warning('Off','all') ;

path(pathdef) ;

addpath('./routines')
addpath('./routines/solmat')

%% Data

setup_model

SET.no_demographics = 0 ;
SET.usecore = 0 ;
SET.unant_demographics = 0 ;

%% Estimation parameters

trend_dates = 1940:0.25:2070 ; 

%convert to quarterly
year_vec  = 1940:2142 ;
datevec_q = 1940:0.25:2142 ;

start_date = 1986 ; SET.start_date=start_date;
end_date   = 2019.75; SET.end_date=end_date;

trend_start = 1960 ;
trend_end   = 2060 ;

trend_dates = trend_dates(find(trend_dates==trend_start):find(trend_dates==trend_end)) ;
SET.trend_dates = trend_dates ;

load_data
data = Zobs;

SET.ss   = length(data(1,:)) ;
nobs     = length(data(:,1)) ;
SET.nobs = nobs ;

SET.EST.hidx      = 1:SET.nobs ;
SET.EST.hidx_no_i = 2:SET.nobs ;

SET.EST.obs_ = [
    SET.variable.iobs
    SET.variable.pinfobs
    SET.variable.dy
    SET.variable.dc
    SET.variable.dinve] ;

% load time-varying adjustment factors from OLG model
load ../olg/adj_factors.mat

% change in life expectancy in years
% take average mortality implied by change in life expectancy
% change the discount factor by the chage in that average mortality

phi_adj = phi_s ;
v_adj   = v_s ;
b_adj   = A_s ;
a_adj   = B_s ;
lgv_adj = lgv_vec ;

% interpolate on annual adjustment factors to get quarterly factors
lgv_adj = [lgv_adj lgv_adj(end)*ones(1,length(year_vec)-length(lgv_adj))];
phi_adj = interp1(year_vec,(1+phi_adj).^1-1,datevec_q,'linear','extrap') ;
v_adj   = interp1(year_vec,(1+v_adj).^1-1,datevec_q,'linear','extrap') ;
b_adj   = interp1(year_vec,(1+b_adj).^1-1,datevec_q,'linear','extrap') ;
a_adj   = interp1(year_vec,(1+a_adj).^1-1,datevec_q,'linear','extrap') ;
tauv    = interp1(year_vec,tauv,datevec_q,'linear','extrap') ;
lgv_adj = interp1(year_vec,(1+lgv_adj).^1-1,datevec_q,'linear','extrap') ;

% restrict changes to a date range
phi_vec = phi_adj(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;
v_vec = v_adj(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;
b_vec = b_adj(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;
a_vec = a_adj(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;
tau_vec = tauv(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;
lgv_vec = lgv_adj(...
    find(datevec_q==trend_start):...
    find(datevec_q==trend_end)) ;

trend_params = ...
    [phi_vec(1) ; v_vec(1) ; b_vec(1) ; a_vec(1) ; tau_vec(1) ; lgv_vec(1)] ;

if SET.no_demographics
    trend_params = ...
        [phi_vec(datevec_q==2000) ; v_vec(datevec_q==2000) ; b_vec(datevec_q==2000) ; a_vec(datevec_q==2000) ; tau_vec(datevec_q==2000) ; lgv_vec(datevec_q==2000)] ;
end
    
save trend_params trend_params

% store vectors for construction of likelihood
SET.tv_str_chg.ant_chg.phi_vec = phi_vec ;
SET.tv_str_chg.ant_chg.v_vec   = v_vec ;
SET.tv_str_chg.ant_chg.b_vec   = b_vec ;
SET.tv_str_chg.ant_chg.a_vec   = a_vec ;
SET.tv_str_chg.ant_chg.tau_vec = tau_vec ;
SET.tv_str_chg.ant_chg.lgv_vec = lgv_vec ;

SET.tv_str_chg.flag = 1 ;

Tbstar = length(SET.tv_str_chg.ant_chg.phi_vec) ;
SET.tv_str_chg.Tbstar = Tbstar ;

SET.tv_str_chg.ant_chg_params   = ...
    {'phi','v','a','b','tau','lgv'} ;
SET.tv_str_chg.unant_chg_params = ...
    {} ; 

gen_set_params(SET,dyn_in) ;

SET.tv_str_chg.cur_t = 1 ;

if SET.no_demographics
    SET.tv_str_chg.flag = 0 ;

    Tbstar = 0 ;
    SET.tv_str_chg.Tbstar = Tbstar ;

    SET.tv_str_chg.ant_chg_params   = ...
        {} ;
    SET.tv_str_chg.unant_chg_params = ...
        {} ; 
end
    
%% fmincon to get hessian

data = data(:,1:find(date_vec==2007.75)) ;

SET.ss = length(data(1,:)) ;

SET.EST.nparam_est = 11 ; 

x0  = M_.params(1:SET.EST.nparam_est) ; 
T_f = zeros(1,length(data)) ;

post = @(x) -ll(x, T_f, SET, data) - prior(x) ;

post(x0)

[llval_csmw,min_param_ll_csmw,gh,hess_csmw,itct,fcount,retcodeh] = ...
     csminwel(post,x0,.01*eye(SET.EST.nparam_est),[],1e-5,20);

hess = inv(hess_csmw) ;
x0   = min_param_ll_csmw ;

%% MCMC

delete(gcp('nocreate'))

SET.nobs = length(data(:,1)) ;
SET.ss   = length(data(1,:)) ; 

load_data
data = Zobs;

SET.ss = length(data(1,:)) ;

SET.EST.chains = 2 ;
SET.EST.kappa = 0.3 ; if SET.no_demographics ; SET.EST.kappa = 0.5; end
SET.EST.N = 200000 ;  if SET.no_demographics ; SET.EST.N = 1000000; end
SET.EST.df = 12 ;

c = SET.EST.kappa * inv(hess) ; % Uses fmincon Hessian
c = chol(c)' ;

SET.EST.chains  = 2 ;
maxproc         = SET.EST.chains;

clear params_x params_T acc_rates
clc

parfor runs=1:maxproc
    warning('Off','all') ;
    
    out = mcmc(SET, data, SET.zlb_t, c, x0, runs) ;
    
    params_x(:,:,runs) = out.xx_out ;
    acc_rates(runs,:)  = out.acc_rates ;
end

delete(gcp('nocreate'))

maxproc = SET.EST.chains;

SET.EST.burn_in = 1/2 ;

params_x_   = [] ;

% Stack chains
for runs=1:maxproc
    params_x_ = [params_x_ ; params_x(round(length(params_x)*SET.EST.burn_in):end,:,runs)] ; 
end

if ~SET.no_demographics
    clear date
    eval(['save mhall_',date,'.mat;']);
end

if SET.no_demographics
    clear date
    eval(['save mhall_',date,'_nodemographics.mat;']);
end    

return

%% Chain diagnostics

sq_R_x = [];
sq_R_T_d = [];
sq_R_T_f = [];

pvec= 1:SET.EST.nparam_est ;

count = 1 ;

yvec = 500:500:length(params_x(:,1,1)) ;

for j=yvec
    j
    sq_R_x(:,count) = ...
        mh_diag(params_x(1:j,pvec,:)) ;
    count = count+1 ;
end

figure ;

plot(yvec,sq_R_x','k','LineWidth',1) ;
line([0 j],[1.1 1.1],'LineWidth',2,'LineStyle','--') ;
line([0 j],[1 1],'LineWidth',1,'color','k') ;

ax = gca ;
ax.YGrid='on';
box('off') ;
set(get(gcf,'CurrentAxes'),'FontName','Times New Roman','FontSize',16,'LineWidth',1.5)
set(gca,'DefaultTextInterpreter', 'latex','TickLabelInterpreter','latex')

xlabel('Length of chain','Interpreter','latex') ;

set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')

%%

date = 1986:0.25:2019.75 ;

figure ;

set(gcf, 'DefaultAxesLineWidth', 1);
set(gcf, 'DefaultLineLineWidth', 1);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',8);

subplot(3,3,1) ;
    plot(date,400*data(1,:))
    line([1986 2020],[0 0],'Color','k','LineWidth',1) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    title('Nominal Fed Funds Rate','Interpreter','latex') ;
    xlim([1986 2020]) ;
subplot(3,3,2) ;
    plot(date,400*data(2,:))
    line([1986 2020],[0 0],'Color','k','LineWidth',1) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    title('Inflation, GDP Deflator','Interpreter','latex') ;   
    xlim([1986 2020]) ;
subplot(3,3,3) ;
    plot(date,100*data(3,:))
    line([1986 2020],[0 0],'Color','k','LineWidth',1) ;
    box('off') ; ax = gca ; ax.YGrid='on';
	title('Real Output Growth','Interpreter','latex') ; 
    xlim([1986 2020]) ;
subplot(3,3,4) ;
    plot(date,100*data(4,:))
    line([1986 2020],[0 0],'Color','k','LineWidth',1) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    title('Real Consumption Growth','Interpreter','latex') ;
    xlim([1986 2020]) ;
subplot(3,3,5) ;
    plot(date,100*data(5,:))
    line([1986 2020],[0 0],'Color','k','LineWidth',1) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    title('Real Investment Growth','Interpreter','latex') ;    
    xlim([1986 2020]) ;
subplot(3,3,6) ;
    bar(date,T_f)
    xlim([2008 2016]) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    title('ZLB Durations, Qtrs','Interpreter','latex') ;    

j = length(params_x_);

figure ;

set(gcf, 'DefaultAxesLineWidth', 1);
set(gcf, 'DefaultLineLineWidth', 1);
set(gcf, 'DefaultAxesTickLabelInterpreter','latex'); 
set(gcf, 'DefaultLegendInterpreter','latex');
set(gcf, 'DefaultAxesFontSize',7);

set(gcf,'PaperUnits','inches','PaperPosition',[0 0 12 5]) ;
subplot(3,4,1) ; hold on ;
    [x,y]=hist(params_x_(1:j,1),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:1,betapdf(0:0.01:1,5.05556,5.05556)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0.25 1]) ;
    title('$\rho_{\chi}$','Interpreter','latex') ;
subplot(3,4,2) ; hold on ;
    [x,y]=hist(params_x_(1:j,2),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:1,betapdf(0:0.01:1,5.05556,5.05556)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0.25 1]) ;
    title('$\rho_{\mu}$','Interpreter','latex') ;
subplot(3,4,3) ; hold on ;
    [x,y]=hist(params_x_(1:j,3),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:1,betapdf(0:0.01:1,5.05556,5.05556)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0.25 1]) ;
    title('$\rho_{\theta}$','Interpreter','latex') ;
subplot(3,4,4) ; hold on ;
    [x,y]=hist(params_x_(1:j,4),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:1,betapdf(0:0.01:1,5.05556,5.05556)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0.25 1]) ;
    title('$\rho_{g}$','Interpreter','latex') ;
subplot(3,4,5) ; hold on ;
    [x,y]=hist(params_x_(1:j,5),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:1,betapdf(0:0.01:1,5.05556,5.05556)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0.25 1]) ;
    title('$\rho_{\kappa}$','Interpreter','latex') ;   
subplot(3,4,6) ; hold on ;
    [x,y]=hist(params_x_(1:j,6),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 4]) ;
    title('$100\times\sigma_\chi$','Interpreter','latex') ;       
subplot(3,4,7) ; hold on ;
    [x,y]=hist(params_x_(1:j,7),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;     
    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 1.5]) ;
    title('$100\times\sigma_\mu$','Interpreter','latex') ;  
subplot(3,4,8) ; hold on ;
    [x,y]=hist(params_x_(1:j,8),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 0.1]) ;
    title('$100\times\sigma_\theta$','Interpreter','latex') ;  
subplot(3,4,9) ; hold on ;
    [x,y]=hist(params_x_(1:j,9),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;       box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 0.2]) ;
    title('$100\times\sigma_R$','Interpreter','latex') ;  
subplot(3,4,10) ; hold on ;
    [x,y]=hist(params_x_(1:j,10),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 4]) ;
    title('$100\times\sigma_g$','Interpreter','latex') ;  
subplot(3,4,11) ; hold on ;
    [x,y]=hist(params_x_(1:j,11),100) ;
    bar(y,x/sum((y(2)-y(1))*x),1.0,'FaceColor','k','EdgeColor','k') ;
    plot(0:0.01:10,unifpdf(0:0.01:10,0,4)) ;    box('off') ; ax = gca ; ax.YGrid='on';
    xlim([0 2]) ;
    title('$100\times\sigma_\kappa$','Interpreter','latex') ;      
    
set(findall(gcf,'-property','FontSize'),'FontSize',12) ;
set(findall(gcf, 'Type', 'Text'),'FontWeight', 'Normal')
